# ==============================================================================
# 文件: initial_trainer.py
# 作用: 实现算法的第一阶段。它接收正样本和未标记样本的嵌入，
#      通过启发式采样方法筛选出可靠负样本（RN），然后训练并返回
#      一个初始的分类器。
# ==============================================================================
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
import numpy as np

import config # 从同级目录导入配置
from models import SimpleMLP # 从同级目录的models模块导入MLP模型

def train_initial_model(p_embeddings, u_embeddings):
    """
    训练初始分类器。
    :param p_embeddings: Tensor, P集的嵌入向量。
    :param u_embeddings: Tensor, U集的嵌入向量。
    :return: 训练好的初始MLP模型。
    """
    print("\n--- 阶段一：开始训练初始模型 ---")
    
    # --- 步骤 1: 筛选可靠负样本 (RN) ---
    print(f"正在筛选可靠负样本 (RN)，使用距离度量: {config.RN_DISTANCE_METRIC}...")
    
    if config.RN_DISTANCE_METRIC == 'centroid_distance':
        # 高效的质心距离法（代码原始实现）
        p_center = p_embeddings.mean(dim=0)
        distances = torch.norm(u_embeddings - p_center, dim=1).numpy()
    elif config.RN_DISTANCE_METRIC == 'mean_distance':
        # 鲁棒的平均距离法（md文档中的公式），计算成本更高
        # 使用广播机制高效计算所有u到所有p的距离矩阵
        # (M, 1, D) - (1, N, D) -> (M, N, D)
        dist_matrix = torch.norm(u_embeddings.unsqueeze(1) - p_embeddings.unsqueeze(0), dim=2)
        # 计算每个u到所有p的平均距离
        distances = dist_matrix.mean(dim=1).numpy()
    else:
        raise ValueError(f"未知的距离度量 '{config.RN_DISTANCE_METRIC}'。请在 'centroid_distance' 或 'mean_distance' 中选择。")

    # 根据分位数计算距离阈值
    d_low = np.quantile(distances, config.RN_QUANTILE_LOW)
    d_high = np.quantile(distances, config.RN_QUANTILE_HIGH)
    
    # 找到距离在指定区间的候选样本的索引
    candidate_indices = np.where((distances >= d_low) & (distances <= d_high))[0]
    
    # 从候选索引中随机选择k个作为RN的索引
    # 使用replace=False确保无放回抽样
    if len(candidate_indices) < config.K_RELIABLE_NEGATIVES:
        print(f"警告: 候选负样本数量({len(candidate_indices)})少于所需的k({config.K_RELIABLE_NEGATIVES})。将使用所有候选样本。")
        rn_indices = candidate_indices
    else:
        rn_indices = np.random.choice(candidate_indices, config.K_RELIABLE_NEGATIVES, replace=False)

    # 获取RN的嵌入向量
    rn_embeddings = u_embeddings[rn_indices]
    print(f"筛选出 {len(rn_embeddings)} 个可靠负样本。")

    # --- 步骤 2: 准备训练数据和模型 ---
    # 创建训练数据和标签
    train_embeddings = torch.cat([p_embeddings, rn_embeddings], dim=0)
    train_labels = torch.cat([
        torch.ones(p_embeddings.shape[0], 1), # P集标签为1
        torch.zeros(rn_embeddings.shape[0], 1) # RN集标签为0
    ], dim=0)
    
    # 创建PyTorch数据集和数据加载器
    initial_dataset = TensorDataset(train_embeddings, train_labels)
    initial_dataloader = DataLoader(initial_dataset, batch_size=config.INITIAL_TRAIN_BATCH_SIZE, shuffle=True)
    
    # 初始化MLP模型并移动到指定设备
    model = SimpleMLP(input_dim=config.EMBEDDING_DIM, hidden_dims=config.HIDDEN_DIMS).to(config.DEVICE)
    
    # 定义损失函数和优化器
    # BCEWithLogitsLoss更数值稳定，因为它将Sigmoid和BCE合并在一起
    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.INITIAL_TRAIN_LR)

    # --- 步骤 3: 训练模型 ---
    print("开始训练初始分类器...")
    model.train() # 将模型设置为训练模式
    for epoch in range(config.INITIAL_TRAIN_EPOCHS):
        total_loss = 0
        for batch_embeds, batch_labels in initial_dataloader:
            # 将数据移动到指定设备
            batch_embeds = batch_embeds.to(config.DEVICE)
            batch_labels = batch_labels.to(config.DEVICE)
            
            # 清零梯度
            optimizer.zero_grad()
            
            # 前向传播，得到logits
            outputs = model(batch_embeds)
            
            # 计算损失
            loss = criterion(outputs, batch_labels)
            
            # 反向传播
            loss.backward()
            
            # 更新参数
            optimizer.step()
            
            total_loss += loss.item()
        
        # 每10轮打印一次训练信息
        if (epoch + 1) % 10 == 0:
            print(f"初始模型训练 - 轮次 [{epoch+1}/{config.INITIAL_TRAIN_EPOCHS}], 平均损失: {total_loss/len(initial_dataloader):.4f}")

    # 保存训练好的初始模型
    print(f"初始模型训练完成，保存到: {config.INITIAL_MODEL_PATH}")
    torch.save(model.state_dict(), config.INITIAL_MODEL_PATH)
    
    return model
